﻿using Microsoft.EntityFrameworkCore.Storage;
using Microsoft.Extensions.Logging;
using Polly;
using System;
using System.Collections.Generic;
using System.Threading.Tasks;
using WorkflowCore.Interface;
using WorkflowCore.Models;

namespace Acme.BookStore
{
    public class PollyRetryStepMiddleware : IWorkflowStepMiddleware
    {
        private const string StepContextKey = "WorkflowStepContext";
        private const int MaxRetries = 3;
        private readonly ILogger<PollyRetryStepMiddleware> _log;

        //public PollyRetryMiddleware(ILogger<PollyRetryStepMiddleware> log)
        public PollyRetryStepMiddleware(ILogger<PollyRetryStepMiddleware> log)
        {
            _log = log;
        }

        // Consult Polly's docs for more information on how to build
        // retry policies:
        // https://github.com/App-vNext/Polly
        public IAsyncPolicy<ExecutionResult> GetRetryPolicy() =>
            Policy<ExecutionResult>
                .Handle<TimeoutException>()
                .RetryAsync(
                    MaxRetries,
                    (result, retryCount, context) =>
                        UpdateRetryCount(
                            result.Exception,
                            retryCount,
                            context[StepContextKey] as IStepExecutionContext)
                );

        public async Task<ExecutionResult> HandleAsync(
            IStepExecutionContext context,
            IStepBody body,
            WorkflowStepDelegate next
        )
        {
            return await GetRetryPolicy().ExecuteAsync(
                ctx => next(),
                // The step execution context gets passed down so that
                // the step is accessible within the retry policy
                new Dictionary<string, object>
                {
                { StepContextKey, context }
                });
        }

        private Task UpdateRetryCount(
            Exception exception,
            int retryCount,
            IStepExecutionContext stepContext)
        {
            var stepInstance = stepContext.ExecutionPointer;
            stepInstance.RetryCount = retryCount;
            return Task.CompletedTask;
        }
    }
}
